import matplotlib.pyplot as plt
import pandas as pd
import glob
import numpy as np

extracted_files = glob.glob("/home/shaonian/SED/SED/data_eng/statistics/output/*.tsv")
extracted_files = [x for x in extracted_files if "val" not in x and "test" not in x]
# extracted_files = [x for x in extracted_files if ("strong_" not in x)]
extracted_files = [x for x in extracted_files if ("od" not in x.lower())]
# draw 4 figures with 4 features
# give me 10 colors
colors = plt.cm.jet(np.linspace(0, 1, len(extracted_files)))
kwargs = dict(histtype='stepfilled', alpha=0.2, density=True, bins=50)

for feat_name in ["spectral_flatness", "pitch_salience", "spectral_complexity", "spectral_flux"]:
    fig = plt.figure(figsize=(8, 8))
    plt.title(f"{feat_name} distribution", fontsize=16)
    
    for i, tsv in enumerate(extracted_files):
        df = pd.read_csv(tsv, sep="\t")
        plt.hist(df[feat_name], **kwargs, color=colors[i])
        label = tsv.split("/")[-1].replace(".tsv", "").replace("train_", "").replace("unlabeled_", "")
        label = "_".join(label.split("_")[:2])
        df[feat_name].plot(kind='kde', color=colors[i], label=label)
    plt.xlabel(feat_name)    
    plt.tight_layout()
    plt.legend(fontsize=16)
    plt.savefig(f"./output/{feat_name}_distribution.png")
    